import json
import numpy as np
import matplotlib.pyplot as plt

# Load JSON data
with open("generative_pm/results_Llama3_length_1shot_mt-banch.json", 'r') as file:
    data = json.load(file)

chosen_score_list = data["chosen_score_list"]

# Convert list to numpy array
scores = np.array(chosen_score_list)

# Calculate absolute distances from 0.5
distances = np.abs(scores - 0.5)

# Define bins with a width of 0.1 between 0 and 0.5
bins = np.arange(0, 0.6, 0.1)

# Bin each score into the corresponding range
bin_indices = np.digitize(distances, bins, right=True) - 1

# Initialize dictionary to store bin results
bin_results = {}

# Iterate through bins
for i in range(len(bins) - 1):
    bin_scores = scores[bin_indices == i]
    Agreement = np.sum(bin_scores > 0.5)
    Disagreement = np.sum(bin_scores <= 0.5)
    bin_results[(bins[i], bins[i + 1])] = {
        'Agreement': Agreement,
        'Disagreement': Disagreement,
        'sample_size': len(bin_scores)
    }

# Prepare data for plotting
bin_ranges = [(f"{k[0]+0.5:.1f}-{k[1]+0.5:.1f}", v['Agreement'], v['Disagreement'], v['sample_size']) for k, v in bin_results.items() if k[0] <= 0.5]

# Create figure and axis
fig, ax1 = plt.subplots(figsize=(8, 6))
bar_width = 0.7
index = np.arange(len(bin_ranges))

# Plot stacked bar chart
bars1 = ax1.bar(index, [x[1] for x in bin_ranges], bar_width, color='#b2df8a', alpha=1, label='Agreement Number')
bars2 = ax1.bar(index, [x[2] for x in bin_ranges], bar_width, bottom=[x[1] for x in bin_ranges], color='#fdaeae', alpha=1, label='Disagreement Number')

# Annotate the stacked bar chart
for i, bar in enumerate(bars1):
    yval = bar.get_height()
    ax1.text(bar.get_x() + bar.get_width() / 2, yval / 2, f"{int(yval)}", ha='center', va='center', fontsize=26)

for i, bar in enumerate(bars2):
    yval = bar.get_height()
    ax1.text(bar.get_x() + bar.get_width() / 2, bars1[i].get_height() + yval / 2, f"{int(yval)}", ha='center', va='center', fontsize=26)

# Remove y-axis label for ax1 (stacked bar chart)
ax1.set_ylabel('')  # Remove the left y-axis label
ax1.tick_params(axis='y', left=False, labelleft=False)  # Hide left y-axis ticks and labels

# Customize x-axis
ax1.set_xlabel('Preference Strength', fontsize=26)
ax1.set_xticks(index)
ax1.set_xticklabels([f"{r[0]}" for r in bin_ranges], fontsize=26)

# Remove top and right spines
ax1.spines['top'].set_visible(False)
ax1.spines['right'].set_visible(False)

# Set tick parameters
ax1.tick_params(axis='both', labelsize=22)

# Add legend for stacked bar chart
ax1.legend(loc='upper left', fontsize=23)

# Create a secondary y-axis for the agreement rate line chart
ax2 = ax1.twinx()

# Calculate Agreement Rate (Agreement / sample_size)
agreement_rates = [(x[1] / x[3]) if x[3] > 0 else 0 for x in bin_ranges]

# Plot the agreement rate as a line chart
ax2.plot(index, agreement_rates, color='black', marker='o', linestyle='-', linewidth=4, markersize=10, label='Agreement Rate')

# Set the label for the secondary y-axis
ax2.set_ylabel('Agreement Rate', fontsize=26, color='black')
ax2.tick_params(axis='y', labelsize=22, colors='black')

# Add the secondary y-axis legend
ax2.legend(loc='upper left', fontsize=23, bbox_to_anchor=(0, 0.75))

# Set y-axis limits for the agreement rate
ax2.set_ylim(0.5, 1)

# Hide the left y-axis line
ax1.spines['left'].set_visible(False)
# Hide the top spine
ax1.spines['top'].set_visible(False)

# Hide the left y-axis line
ax2.spines['left'].set_visible(False)
# Hide the top spine
ax2.spines['top'].set_visible(False)

# Hide the left y-axis ticks
ax1.tick_params(axis='y', which='both', left=False)


# Save and show the plot
plt.tight_layout()
plt.savefig('generative_pm/sample_size_bins-Llama3_length1shot_MTbench_stacked_with_agreement_rate_no_number_axis.png')
plt.show()
